#!/usb/bin/python

import cStringIO
import random
import optparse

import sys
from boto import ec2
from os import path

def main ():

    parser = optparse.OptionParser()
    parser.add_option('--key', '-k', dest='key', default='', 
                      action='store',
                      help='Amazon EC2 private key location')

    parser.add_option('--exta_config', dest='extra_config', 
                      default='~/.cheops_extra_config.py', 
                      action='store',
                      help='Extra configuration file location')

    parser.add_option('--access_id', dest='access_id', default='', 
                      action='store',
                      help='Access ID for S3 File system unit tests')

    parser.add_option('--access_key', dest='access_key', default='',
                      action='store',
                      help='Access Key for S3 File system unit tests')

    parser.add_option('--reservation', dest='reservation', default=None,
                      action='store', type='int',
                      help='Reservation number (all if not specified)')

    (options, args) = parser.parse_args()

    extra_config = path.expanduser(options.extra_config)
    if options.reservation == None:
        reservation_id = 'all'
    else:
        reservation_id = 'rev%d' % (options.reservation)
    f_config = open('cheops_config.%s.py' % (reservation_id), 'w')
    f_ssh = open('run_cheops.%s.sh' % (reservation_id), 'w') 
    f_exec = open('update_execute.%s.sh' % (reservation_id), 'w') 
    f_uconfig = open('update_config.%s.sh' % (reservation_id), 'w') 
    port = random.randint(10000, 50000)

    f_config.write("import math\n")

    ec2_conn = ec2.connection.EC2Connection(
        options.access_id, options.access_key)
    reservations = ec2_conn.get_all_instances()
    if options.reservation != None:
        try:
            reservations = [reservations[options.reservation]]
        except IndexError:
            print >>sys.stderr, (
                'There is no reservation with index %d' % (
                    options.reservation))

    for reservation in reservations:
        instances = reservation.instances
        instances.sort(key=lambda x: int(x.ami_launch_index))

        for instance in [i for i in instances if i.state == 'running']:
            if int(instance.ami_launch_index) == 0:
                f_config.write('DISPATCHER = (\'%s\', %d)\n' % (
                        instance.private_ip_address, port))
            else:
                f_config.write('HOSTS[%d] = (\'%s\', %d)\n' % (
                        int(instance.ami_launch_index),
                        instance.private_ip_address, port))
            f_ssh.write('ssh -o "StrictHostKeyChecking no" '
                        'ubuntu@%s -i %s '
                        'source execute.sh' % (
                    instance.public_dns_name, options.key))
            if int(instance.ami_launch_index) != 0:
                f_ssh.write(' & \n')
            else:
                f_ssh.write('\nsleep 10\n')
            f_exec.write('scp -o "StrictHostKeyChecking no" '
                        '-i %s '
                        './execute.sh ubuntu@%s:./execute.sh &\n' % (
                    options.key, instance.public_dns_name))
            f_uconfig.write('scp -o "StrictHostKeyChecking no" '
                            '-i %s '
                            './cheops_config.%s.py '
                            'ubuntu@%s:./.cheops.py &\n' % (
                    options.key, reservation_id, instance.public_dns_name))
            
                
    if path.isfile(extra_config):
        f_config.write(open(extra_config).read())
    f_config.write('\n')
    f_config.close()
    f_ssh.close()
    f_exec.close()

if __name__ == '__main__':
    main ()
